#include "Basis.h"
#include "Memory.h"

PMDL AllocateWritableMdl(ULONG virtualAddress, ULONG length, PULONG pNewVirtualAddress)
{
	PMDL pMdl = IoAllocateMdl((PVOID)virtualAddress, length, FALSE, FALSE, NULL);
	if (!pMdl)
	{
		KdPrint(("failed to allocate mdl for 0x%x", virtualAddress));
		return NULL;
	}

	MmProbeAndLockPages(pMdl, KernelMode, IoWriteAccess);
	if (pMdl->MdlFlags & (MDL_MAPPED_TO_SYSTEM_VA | MDL_SOURCE_IS_NONPAGED_POOL))
	{
		*pNewVirtualAddress = (ULONG)pMdl->MappedSystemVa;
	}
	else
	{
		*pNewVirtualAddress = (ULONG)MmMapLockedPagesSpecifyCache(pMdl, KernelMode, MmCached, NULL, FALSE, NormalPagePriority);
	}

	if (*pNewVirtualAddress == 0)
	{
		KdPrint(("failed to make writable address for 0x%x", virtualAddress));
		FreeMdl(pMdl);
		pMdl = NULL;
	}

	return pMdl;
}

VOID FreeMdl(PMDL pMdl)
{
	MmUnlockPages(pMdl);
	IoFreeMdl(pMdl);
}

VOID FreeList(PLIST_ENTRY pListEntry, PKSPIN_LOCK pSpinLock, PNPAGED_LOOKASIDE_LIST pLookaside, FreeListItemCallback freeListItemCallback)
{
	KIRQL oldIrq;
	KeAcquireSpinLock(pSpinLock, &oldIrq);
	while (!IsListEmpty(pListEntry))
	{
		PLIST_ENTRY pItem = RemoveHeadList(pListEntry);
		if (freeListItemCallback != NULL)
		{
			freeListItemCallback(pItem);
		}

		if (pLookaside != NULL)
		{
			ExFreeToNPagedLookasideList(pLookaside, pItem); // free in lookaside
		}
		else
		{
			ExFreePool(pItem);
		}
	}
	KeReleaseSpinLock(pSpinLock, oldIrq);
}

NDIS_STATUS CopyNblToBuffer(IN NDIS_HANDLE bindingHandle, PNET_BUFFER_LIST netBufferList, PVOID *ppBuffer, PULONG pBytesCopied)
{
	NDIS_STATUS status = NDIS_STATUS_SUCCESS;
	*pBytesCopied = 0;

	__try
	{
		PNET_BUFFER nextNetBuffer;
		for (nextNetBuffer = NET_BUFFER_LIST_FIRST_NB(netBufferList); nextNetBuffer != NULL; nextNetBuffer = NET_BUFFER_NEXT_NB(nextNetBuffer))
		{
			BOOLEAN isFirstMdl = TRUE;
			PMDL nextMdl = NET_BUFFER_CURRENT_MDL(nextNetBuffer);
			while (nextMdl)
			{
				PUCHAR pData = NULL;
				ULONG dataLength = 0;
				NdisQueryMdl(nextMdl, &pData, &dataLength, NormalPagePriority);
				if (pData == NULL)
				{
					// failed to get source
					if (*ppBuffer && *pBytesCopied > 0)
					{
						NdisFreeMemory(*ppBuffer, 0, 0);
						*ppBuffer = NULL;
					}
					status = NDIS_STATUS_RESOURCES;
					__leave;
				}

				if (isFirstMdl)
				{
					pData += NET_BUFFER_CURRENT_MDL_OFFSET(nextNetBuffer);
					dataLength -= NET_BUFFER_CURRENT_MDL_OFFSET(nextNetBuffer);
					isFirstMdl = FALSE;
				}

				if (dataLength > 0)
				{
					ULONG newLength = *pBytesCopied + dataLength;
					PUCHAR newBuffer = (PUCHAR)NdisAllocateMemoryWithTagPriority(bindingHandle, newLength, NETWORK_ALLOC_TAG, NormalPoolPriority);
					if (newBuffer == NULL)
					{
						// failed to allocate memory
						if (*ppBuffer && *pBytesCopied > 0)
						{
							NdisFreeMemory(*ppBuffer, 0, 0);
							*ppBuffer = NULL;
						}
						status = NDIS_STATUS_RESOURCES;
						__leave;
					}

					// copy existing data to new buffer and free existing buffer
					if (*ppBuffer && *pBytesCopied > 0)
					{
						NdisMoveMemory(newBuffer, *ppBuffer, *pBytesCopied);
						ExFreePool(*ppBuffer);
					}

					// copy new data
					NdisMoveMemory(newBuffer + *pBytesCopied, pData, dataLength);

					// set new length and buffer
					*pBytesCopied = newLength;
					*ppBuffer = newBuffer;
				}

				// next
				nextMdl = NDIS_MDL_LINKAGE(nextMdl);
			}
		}
	}
	finally {}

	return status;
}
